//
// Created by mio on 25-5-23.
//

#include "DXGICapture.h"

// FFmpeg头文件，如果AVFramePtr等没有完全包含，确保包含
extern "C" {
#include <libavutil/imgutils.h> // For av_image_fill_arrays if needed directly, though AVFrame manages its buffers
}

DXGICapture::DXGICapture()
    : d3d_device_(nullptr),
      d3d_context_(nullptr),
      desk_dupl_(nullptr),
      staging_texture_(nullptr),
      capture_full_screen_(true)
{
    ZeroMemory(&texture_desc_, sizeof(texture_desc_));
    ZeroMemory(&capture_rect_, sizeof(capture_rect_));
}

DXGICapture::~DXGICapture()
{
    Close();
}

bool DXGICapture::GetDisplayRect(qint32 display_index, RECT& display_rect)
{
    IDXGIFactory1* dxgi_factory = nullptr;
    HRESULT hr = CreateDXGIFactory1(__uuidof(IDXGIFactory1), (void**)&dxgi_factory);
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to create DXGI Factory:" << hr;
        return false;
    }

    IDXGIAdapter1* dxgi_adapter = nullptr;
    hr = dxgi_factory->EnumAdapters1(0, &dxgi_adapter); // 通常使用第一个适配器
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to enumerate DXGI Adapters:" << hr;
        if (dxgi_factory) dxgi_factory->Release();
        return false;
    }

    IDXGIOutput* dxgi_output = nullptr;
    hr = dxgi_adapter->EnumOutputs(display_index, &dxgi_output);
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to enumerate DXGI Outputs for display index" << display_index << ":" << hr;
        if (dxgi_adapter) dxgi_adapter->Release();
        if (dxgi_factory) dxgi_factory->Release();
        return false;
    }

    DXGI_OUTPUT_DESC output_desc;
    hr = dxgi_output->GetDesc(&output_desc);
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to get DXGI Output Description:" << hr;
    } else {
        display_rect = output_desc.DesktopCoordinates;
    }

    if (dxgi_output) dxgi_output->Release();
    if (dxgi_adapter) dxgi_adapter->Release();
    if (dxgi_factory) dxgi_factory->Release();
    return SUCCEEDED(hr);
}


bool DXGICapture::InitializeDXGI(qint32 display_index, RECT target_capture_rect) {
    HRESULT hr;

    // 1. 创建D3D设备和上下文
    D3D_FEATURE_LEVEL feature_levels[] = { D3D_FEATURE_LEVEL_11_0, D3D_FEATURE_LEVEL_10_1 };
    hr = D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0,
                           feature_levels, ARRAYSIZE(feature_levels), D3D11_SDK_VERSION,
                           &d3d_device_, nullptr, &d3d_context_);
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to create D3D11 device:" << hr;
        return false;
    }

    // 2. 获取DXGI设备
    IDXGIDevice* dxgi_device = nullptr;
    hr = d3d_device_->QueryInterface(__uuidof(IDXGIDevice), reinterpret_cast<void**>(&dxgi_device));
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to query IDXGIDevice:" << hr;
        return false;
    }

    // 3. 获取DXGI适配器
    IDXGIAdapter* dxgi_adapter = nullptr;
    hr = dxgi_device->GetParent(__uuidof(IDXGIAdapter), reinterpret_cast<void**>(&dxgi_adapter));
    dxgi_device->Release(); // dxgi_device 不再需要
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to get IDXGIAdapter:" << hr;
        return false;
    }

    // 4. 获取DXGI输出 (显示器)
    IDXGIOutput* dxgi_output = nullptr;
    hr = dxgi_adapter->EnumOutputs(display_index, &dxgi_output);
    dxgi_adapter->Release(); // dxgi_adapter 不再需要
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to enumerate DXGIOutput for display index" << display_index << ":" << hr;
        return false;
    }

    DXGI_OUTPUT_DESC output_desc;
    hr = dxgi_output->GetDesc(&output_desc);
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to get output description:" << hr;
        dxgi_output->Release();
        return false;
    }

    // 设置捕获区域
    if (capture_full_screen_) {
        capture_rect_ = output_desc.DesktopCoordinates;
    } else {
        // 确保用户指定的区域在显示器范围内，并进行裁剪（如果需要）
        // 这里简化为直接使用，但实际应用可能需要更复杂的逻辑来处理多显示器和区域指定
        capture_rect_ = target_capture_rect;
        // 验证 capture_rect_ 是否在 output_desc.DesktopCoordinates 内，并调整
        if (capture_rect_.left < output_desc.DesktopCoordinates.left) capture_rect_.left = output_desc.DesktopCoordinates.left;
        if (capture_rect_.top < output_desc.DesktopCoordinates.top) capture_rect_.top = output_desc.DesktopCoordinates.top;
        if (capture_rect_.right > output_desc.DesktopCoordinates.right) capture_rect_.right = output_desc.DesktopCoordinates.right;
        if (capture_rect_.bottom > output_desc.DesktopCoordinates.bottom) capture_rect_.bottom = output_desc.DesktopCoordinates.bottom;
    }

    width_ = capture_rect_.right - capture_rect_.left;
    height_ = capture_rect_.bottom - capture_rect_.top;

    if (width_ == 0 || height_ == 0) {
        qDebug() << "DXGICapture: Invalid capture dimensions (width or height is 0).";
        dxgi_output->Release();
        return false;
    }


    // 5. 获取IDXGIOutput1接口以进行桌面复制
    IDXGIOutput1* dxgi_output1 = nullptr;
    hr = dxgi_output->QueryInterface(__uuidof(IDXGIOutput1), reinterpret_cast<void**>(&dxgi_output1));
    dxgi_output->Release(); // dxgi_output 不再需要
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to query IDXGIOutput1:" << hr;
        return false;
    }

    // 6. 创建桌面复制接口
    hr = dxgi_output1->DuplicateOutput(d3d_device_, &desk_dupl_);
    dxgi_output1->Release(); // dxgi_output1 不再需要
    if (FAILED(hr)) {
        qDebug() << "DXGICapture: Failed to duplicate output:" << hr;
        // 可能的原因：需要管理员权限，或者其他应用正在独占（如远程桌面，全屏游戏保护等）
        // 或者达到了最大并发复制数
        return false;
    }

    // 7. 获取捕获到的纹理描述 (在第一次捕获帧时获取更准确)
    // 这里可以预先设置一个大致的 texture_desc_ for staging texture
    ZeroMemory(&texture_desc_, sizeof(texture_desc_));
    texture_desc_.Width = width_;
    texture_desc_.Height = height_;
    texture_desc_.Format = DXGI_FORMAT_B8G8R8A8_UNORM; // GDI兼容的通常是BGRA
    texture_desc_.ArraySize = 1;
    texture_desc_.MipLevels = 1;
    texture_desc_.SampleDesc.Count = 1;
    texture_desc_.Usage = D3D11_USAGE_STAGING;
    texture_desc_.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    texture_desc_.BindFlags = 0;
    texture_desc_.MiscFlags = 0;


    qDebug() << "DXGICapture: DXGI Initialized. Capturing area:"
             << capture_rect_.left << capture_rect_.top
             << width_ << "x" << height_;

    return true;
}


bool DXGICapture::Init(qint32 display_index, const QRect& capture_region)
{
    if (is_initialized_) {
        qDebug() << "DXGICapture: Already initialized.";
        return true;
    }
    std::lock_guard<std::mutex> lock(mutex_); // 保护初始化过程

    RECT target_rect = {0,0,0,0};
    if (capture_region.width() > 0 && capture_region.height() > 0) {
        capture_full_screen_ = false;
        target_rect.left = capture_region.left();
        target_rect.top = capture_region.top();
        target_rect.right = capture_region.left() + capture_region.width();
        target_rect.bottom = capture_region.top() + capture_region.height();
    } else {
        capture_full_screen_ = true;
        // 如果是全屏，target_rect 会在 InitializeDXGI 中根据显示器实际尺寸填充
    }


    if (!InitializeDXGI(display_index, target_rect)) {
        Cleanup();
        return false;
    }

    is_initialized_ = true;
    stop_flag_ = false;
    this->start(); // 启动采集线程 (QThread::start()调用run())
    return true;
}

void DXGICapture::Cleanup()
{
    if (staging_texture_) {
        staging_texture_->Release();
        staging_texture_ = nullptr;
    }
    if (desk_dupl_) {
        desk_dupl_->Release();
        desk_dupl_ = nullptr;
    }
    if (d3d_context_) {
        d3d_context_->Release();
        d3d_context_ = nullptr;
    }
    if (d3d_device_) {
        d3d_device_->Release();
        d3d_device_ = nullptr;
    }
    width_ = 0;
    height_ = 0;
    is_initialized_ = false;
    qDebug() << "DXGICapture: Resources cleaned up.";
}

bool DXGICapture::Close()
{
    if (!is_initialized_ && !this->isRunning()) { // 如果未初始化且线程未运行
        return true;
    }

    stop_flag_ = true;
    if (this->isRunning()) {
        this->quit(); // 请求线程退出事件循环 (如果run中有事件循环)
        if (!this->wait(3000)) { // 等待线程结束，超时3秒
            qDebug() << "DXGICapture: Thread wait timeout, terminating...";
            this->terminate(); // 强制终止 (不推荐，但作为后备)
            this->wait();      // 等待终止完成
        }
    }

    std::lock_guard<std::mutex> lock(mutex_); // 保护清理过程
    Cleanup();

    // 清理帧数据
    std::lock_guard<std::mutex> frame_lock(frame_mutex_);
    captured_bgra_frame_.reset();

    return true;
}


void DXGICapture::run() {
    qDebug() << "DXGICapture: Thread started. Target FPS:" << target_framerate_;
    if (!is_initialized_) {
        qDebug() << "DXGICapture: Not initialized, thread exiting.";
        return;
    }

    HRESULT hr;
    IDXGIResource* desktop_resource = nullptr;
    DXGI_OUTDUPL_FRAME_INFO frame_info;
    ID3D11Texture2D* acquired_texture = nullptr;

    // 为 AVFrame 分配内存
    AVFramePtr current_frame(av_frame_alloc(), [](AVFrame* ptr){ av_frame_free(&ptr); });
    if (!current_frame) {
        qDebug() << "DXGICapture: Failed to allocate AVFrame.";
        return;
    }
    current_frame->width = width_;
    current_frame->height = height_;
    current_frame->format = AV_PIX_FMT_BGRA; // DXGI桌面复制通常是BGRA
    if (av_frame_get_buffer(current_frame.get(), 0) < 0) {
        qDebug() << "DXGICapture: Failed to get buffer for AVFrame.";
        return;
    }


    while (!stop_flag_) {
        auto frame_start_time = std::chrono::steady_clock::now();

        // 释放上一帧的资源（如果存在）
        if (acquired_texture) {
            acquired_texture->Release();
            acquired_texture = nullptr;
        }
        if (desktop_resource) {
            desktop_resource->Release();
            desktop_resource = nullptr;
        }
        if (desk_dupl_) { // 每次循环前尝试释放上一帧，避免锁死
             // desk_dupl_->ReleaseFrame(); // 应该在获取下一帧之前或之后立即释放
        }


        // 1. 获取下一帧
        hr = desk_dupl_->AcquireNextFrame(0, &frame_info, &desktop_resource); // 50ms超时

        if (hr == DXGI_ERROR_WAIT_TIMEOUT) {
            //qDebug() << "DXGICapture: AcquireNextFrame timeout (no screen update).";
            // 没有屏幕更新，不是错误，继续尝试
            std::this_thread::sleep_for(std::chrono::milliseconds(1000 / target_framerate_ / 2)); // 短暂休眠
            continue;
        }
        if (FAILED(hr)) {
            qDebug() << "DXGICapture: Failed to acquire next frame:" << hr;
            // 尝试重新初始化DXGI？或者直接退出？
            // 如果是 DXGI_ERROR_ACCESS_LOST，需要重新初始化 desk_dupl_
            if (hr == DXGI_ERROR_ACCESS_LOST) {
                qDebug() << "DXGICapture: DXGI_ERROR_ACCESS_LOST, attempting to reinitialize...";
                // 简单处理：退出线程，让外部重新Init
                Close(); // 会设置 stop_flag_ 和 is_initialized_
                // 或者在这里尝试重新创建 desk_dupl_，会更复杂
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(100)); // 发生错误，稍作等待
            continue;
        }

        // 2. 查询纹理接口
        hr = desktop_resource->QueryInterface(__uuidof(ID3D11Texture2D), reinterpret_cast<void**>(&acquired_texture));
        if (FAILED(hr)) {
            qDebug() << "DXGICapture: Failed to query ID3D11Texture2D from desktop resource:" << hr;
            desk_dupl_->ReleaseFrame(); // 必须释放
            if(desktop_resource) { desktop_resource->Release(); desktop_resource = nullptr; }
            continue;
        }

        // 获取实际捕获纹理的描述，可能与我们预设的 staging_texture_ 不同
        D3D11_TEXTURE2D_DESC current_texture_desc;
        acquired_texture->GetDesc(&current_texture_desc);

        // 3. 准备 Staging Texture
        // 如果 staging_texture_ 不存在，或者尺寸/格式不匹配，则重新创建
        // (考虑到区域截图，这里staging_texture_的尺寸应该基于我们指定的width_和height_，而不是current_texture_desc)
        if (!staging_texture_ ||
            texture_desc_.Width != width_ ||
            texture_desc_.Height != height_ ||
            texture_desc_.Format != current_texture_desc.Format) // 格式通常是固定的BGRA8
        {
            if (staging_texture_) staging_texture_->Release();

            texture_desc_.Width = width_;
            texture_desc_.Height = height_;
            texture_desc_.Format = current_texture_desc.Format; // 使用捕获到的实际格式
            texture_desc_.ArraySize = 1;
            texture_desc_.MipLevels = 1;
            texture_desc_.SampleDesc.Count = 1;
            texture_desc_.SampleDesc.Quality = 0;
            texture_desc_.Usage = D3D11_USAGE_STAGING;
            texture_desc_.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
            texture_desc_.BindFlags = 0;
            texture_desc_.MiscFlags = 0;

            hr = d3d_device_->CreateTexture2D(&texture_desc_, nullptr, &staging_texture_);
            if (FAILED(hr)) {
                qDebug() << "DXGICapture: Failed to create staging texture:" << hr;
                staging_texture_ = nullptr; // 确保置空
                desk_dupl_->ReleaseFrame();
                if(desktop_resource) { desktop_resource->Release(); desktop_resource = nullptr; }
                if(acquired_texture) { acquired_texture->Release(); acquired_texture = nullptr; }
                continue;
            }
            qDebug() << "DXGICapture: Staging texture created/recreated. Format:" << texture_desc_.Format;
        }

        // 4. 复制资源到 Staging Texture
        // 如果是全屏捕获，直接复制
        // 如果是区域捕获，需要使用 CopySubresourceRegion
        if (capture_full_screen_ ||
            (current_texture_desc.Width == width_ && current_texture_desc.Height == height_)) {
            d3d_context_->CopyResource(staging_texture_, acquired_texture);
        } else {
            D3D11_BOX source_region;
            source_region.left = capture_rect_.left - (current_texture_desc.Width == GetSystemMetrics(SM_CXVIRTUALSCREEN) ? GetSystemMetrics(SM_XVIRTUALSCREEN) : 0); // 调整为相对源纹理的坐标
            source_region.top = capture_rect_.top - (current_texture_desc.Height == GetSystemMetrics(SM_CYVIRTUALSCREEN) ? GetSystemMetrics(SM_YVIRTUALSCREEN) : 0);
            source_region.right = source_region.left + width_;
            source_region.bottom = source_region.top + height_;
            source_region.front = 0;
            source_region.back = 1;

            // 确保源区域不超出实际捕获纹理的边界
            if(source_region.left < 0) source_region.left = 0;
            if(source_region.top < 0) source_region.top = 0;
            if(source_region.right > current_texture_desc.Width) source_region.right = current_texture_desc.Width;
            if(source_region.bottom > current_texture_desc.Height) source_region.bottom = current_texture_desc.Height;

            if (source_region.right > source_region.left && source_region.bottom > source_region.top) {
                 d3d_context_->CopySubresourceRegion(staging_texture_, 0, 0, 0, 0, acquired_texture, 0, &source_region);
            } else {
                 qDebug() << "DXGICapture: Invalid source region for CopySubresourceRegion.";
                 // 可以选择清空 staging_texture_ 或跳过此帧
                 desk_dupl_->ReleaseFrame();
                 if(desktop_resource) { desktop_resource->Release(); desktop_resource = nullptr; }
                 if(acquired_texture) { acquired_texture->Release(); acquired_texture = nullptr; }
                 continue;
            }
        }


        // 5. Map Staging Texture 获取数据
        D3D11_MAPPED_SUBRESOURCE mapped_resource;
        hr = d3d_context_->Map(staging_texture_, 0, D3D11_MAP_READ, 0, &mapped_resource);
        if (FAILED(hr)) {
            qDebug() << "DXGICapture: Failed to map staging texture:" << hr;
            desk_dupl_->ReleaseFrame();
            if(desktop_resource) { desktop_resource->Release(); desktop_resource = nullptr; }
            if(acquired_texture) { acquired_texture->Release(); acquired_texture = nullptr; }
            continue;
        }

        // 6. 填充 AVFrame (BGRA格式)
        // current_frame 已经在循环外初始化并分配了buffer
        // 确保 current_frame 的尺寸和格式与 staging_texture_ 匹配
        if (current_frame->width != (int)width_ ||
            current_frame->height != (int)height_ ||
            current_frame->format != AV_PIX_FMT_BGRA) // 假设 staging texture 是 BGRA
        {
            // 尺寸或格式不匹配，需要重新分配 AVFrame 缓冲区
            // （理论上初始化后尺寸固定，除非 Init 重新调用）
            // 这里简单处理，如果发生则报错并跳过。Init应确保尺寸固定。
            qDebug() << "DXGICapture: AVFrame dimensions/format mismatch with staging texture.";
        } else {
            uint8_t* src_data = static_cast<uint8_t*>(mapped_resource.pData);
            int src_linesize = mapped_resource.RowPitch;
            uint8_t* dst_data = current_frame->data[0];
            int dst_linesize = current_frame->linesize[0];

            // 拷贝数据 (通常 BGRA 是紧凑的，dst_linesize == width * 4)
            // 如果 src_linesize == dst_linesize，可以直接拷贝整个区域
            if (src_linesize == dst_linesize && src_linesize == (int)width_ * 4) {
                memcpy(dst_data, src_data, (size_t)height_ * src_linesize);
            } else {
                // 逐行拷贝以处理可能的行间距差异
                for (uint32_t i = 0; i < height_; ++i) {
                    memcpy(dst_data + (size_t)i * dst_linesize,
                           src_data + (size_t)i * src_linesize,
                           (size_t)width_ * 4); // BGRA 每像素4字节
                }
            }
            current_frame->pts = std::chrono::duration_cast<std::chrono::milliseconds>(
                                 std::chrono::steady_clock::now().time_since_epoch()).count(); // 简单PTS

            // 将捕获到的帧放入共享指针，供外部获取
            {
                std::lock_guard<std::mutex> frame_lock(frame_mutex_);
                // 为了创建新的AVFramePtr，我们需要复制current_frame的内容到一个新的AVFrame对象
                // 或者外部直接使用 current_frame 的拷贝，但要保证线程安全
                // 这里我们让 captured_bgra_frame_ 指向 current_frame，外部获取时拷贝
                // 或者更好的做法是，外部 CaptureFrame 也返回一个新拷贝的 AVFramePtr
                // 为简单起见，这里直接替换
                if (!captured_bgra_frame_ || captured_bgra_frame_.use_count() == 1) { // 如果没有被外部持有或只有自己持有
                    captured_bgra_frame_ = current_frame; // 让 captured_bgra_frame_ 接管 current_frame 的管理
                                                         // 并重新为 current_frame 分配，为下一轮做准备
                    current_frame.reset(av_frame_alloc(), [](AVFrame* ptr){ av_frame_free(&ptr); });
                    if (current_frame) {
                        current_frame->width = width_;
                        current_frame->height = height_;
                        current_frame->format = AV_PIX_FMT_BGRA;
                        av_frame_get_buffer(current_frame.get(), 0); // 忽略错误检查以简化
                    } else {
                        qDebug() << "DXGICapture: Failed to re-allocate AVFrame for next capture.";
                        stop_flag_ = true; // 严重错误，停止
                    }
                } else {
                    // captured_bgra_frame_ 仍被外部使用，我们需要创建一个新的 AVFrame 来存放当前数据
                    AVFramePtr new_captured_frame(av_frame_alloc(), [](AVFrame* ptr){ av_frame_free(&ptr); });
                    if (new_captured_frame) {
                        av_frame_copy_props(new_captured_frame.get(), current_frame.get());
                        new_captured_frame->width = current_frame->width;
                        new_captured_frame->height = current_frame->height;
                        new_captured_frame->format = current_frame->format;
                        if (av_frame_get_buffer(new_captured_frame.get(), 0) == 0) {
                            av_image_copy(new_captured_frame->data, new_captured_frame->linesize,
                                          (const uint8_t**)current_frame->data, current_frame->linesize,
                                          (AVPixelFormat)current_frame->format, current_frame->width, current_frame->height);
                            captured_bgra_frame_ = new_captured_frame;
                        }
                    }
                }
            }

            // FPS 统计
            LogFPS();
        }

        // 7. Unmap Staging Texture
        d3d_context_->Unmap(staging_texture_, 0);

        // 8. 释放帧
        hr = desk_dupl_->ReleaseFrame();
        if (FAILED(hr)) {
             qDebug() << "DXGICapture: Failed to release frame:" << hr;
             // 可能需要重新初始化
        }


        // 帧率控制
        // auto frame_end_time = std::chrono::steady_clock::now();
        // auto frame_duration = std::chrono::duration_cast<std::chrono::milliseconds>(frame_end_time - frame_start_time);
        // qDebug() << "Frame duration: :" << frame_duration;
        // long long sleep_duration_ms = (1000 / target_framerate_) - frame_duration.count();
        // qDebug() << "Sleep duration: :" << sleep_duration_ms;
        // if (sleep_duration_ms > 0) {
        //     std::this_thread::sleep_for(std::chrono::milliseconds(sleep_duration_ms));
        // }

    }

    // 循环结束后，清理资源 (acquired_texture 和 desktop_resource 可能在最后一次循环中未释放)
    if (acquired_texture) acquired_texture->Release();
    if (desktop_resource) desktop_resource->Release();

    // 确保 current_frame 被正确释放 (通过 shared_ptr 的析构)

    qDebug() << "DXGICapture: Thread finished.";
}


bool DXGICapture::CaptureFrame(AVFramePtr& bgra_frame_out)
{
    if (stop_flag_ || !is_initialized_) {
        bgra_frame_out.reset();
        return false;
    }

    std::lock_guard<std::mutex> lock(frame_mutex_);
    if (!captured_bgra_frame_ || captured_bgra_frame_->width == 0) { // 检查帧是否有效
        bgra_frame_out.reset();
        return false;
    }

    // 返回捕获帧的拷贝，以避免多线程问题和数据被覆盖
    // （或者约定外部快速处理并释放，但拷贝更安全）
    bgra_frame_out.reset(av_frame_clone(captured_bgra_frame_.get()), [](AVFrame* ptr){ av_frame_free(&ptr); });
    if (!bgra_frame_out) {
        qDebug() << "DXGICapture: Failed to clone captured frame.";
        return false;
    }
    return true;
}
